from PIL import Image
import numpy as np


def keep_text_color(input_path, output_path, text_color_range=None, tolerance=20):
    """
    保留文字颜色，其他区域变为黑色

    参数:
        input_path: 输入图片路径
        output_path: 输出图片路径
        text_color_range: 文字颜色范围
                          格式1: 单颜色 ((r, g, b),) 配合tolerance使用
                          格式2: 颜色区间 ((r_min, g_min, b_min), (r_max, g_max, b_max))
                          如果为None，默认使用浅棕色范围((150, 120, 80), (200, 170, 120))
        tolerance: 颜色容差，当text_color_range为单颜色时有效，范围0-255
    """
    # 打开图片并转换为RGB模式
    with Image.open(input_path) as img:
        # 处理透明图片
        if img.mode == 'RGBA':
            # 分离alpha通道
            r, g, b, a = img.split()
            img = Image.merge('RGB', (r, g, b))
            has_alpha = True
        else:
            img = img.convert('RGB')
            has_alpha = False

        img_array = np.array(img)

        # 默认文字颜色范围（浅棕色系，可根据实际情况调整）
        if text_color_range is None:
            text_color_range = ((150, 120, 80), (200, 170, 120))

        # 解析颜色范围
        if len(text_color_range) == 1:
            # 单颜色 + 容差模式
            r_center, g_center, b_center = text_color_range[0]
            r_min, r_max = max(0, r_center - tolerance), min(255, r_center + tolerance)
            g_min, g_max = max(0, g_center - tolerance), min(255, g_center + tolerance)
            b_min, b_max = max(0, b_center - tolerance), min(255, b_center + tolerance)
        else:
            # 颜色区间模式
            (r_min, g_min, b_min), (r_max, g_max, b_max) = text_color_range
            # 确保颜色值在0-255范围内
            r_min, r_max = max(0, r_min), min(255, r_max)
            g_min, g_max = max(0, g_min), min(255, g_max)
            b_min, b_max = max(0, b_min), min(255, b_max)

        # 创建掩码：在文字颜色范围内的像素保留
        mask = (
                (img_array[:, :, 0] >= r_min) & (img_array[:, :, 0] <= r_max) &
                (img_array[:, :, 1] >= g_min) & (img_array[:, :, 1] <= g_max) &
                (img_array[:, :, 2] >= b_min) & (img_array[:, :, 2] <= b_max)
        )

        # 将不在文字颜色范围内的像素设为黑色
        img_array[~mask] = [0, 0, 0]

        # 恢复alpha通道（如果存在）
        result_img = Image.fromarray(img_array)
        if has_alpha:
            result_img = Image.merge('RGBA', (*result_img.split(), a))

        # 保存处理后的图片
        result_img.save(output_path)

    print(f"处理完成，已保存至{output_path}")
    return output_path


# 使用示例
if __name__ == "__main__":
    # 输入图片路径
    i = 3
    input_image = f"./img/{i}.jpg"
    # 输出图片路径
    output_image = f"./to_img/{i}.jpg"

    # 示例2：使用单颜色+容差模式
    single_color = ((150, 130, 100),)  # 注意元组的格式
    keep_text_color(
        input_path=input_image,
        output_path=output_image,
        text_color_range=single_color,
        tolerance=34
    )
